import numpy as np
import scipy as sp
from scipy.linalg import lstsq

def iterative_refinement_solve(V, L_rho, max_iter=15, tolerance=1e-15, verbose=True):
    """
    迭代精化求解复数线性系统 S @ a = b，其中 S = V^H @ V, b = V^H @ L_rho
    
    Parameters:
    -----------
    V : ndarray, shape (m, n)
        输入复数矩阵
    L_rho : ndarray, shape (m,)
        右端复数向量
    max_iter : int
        最大迭代次数
    tolerance : float
        收敛容差
    verbose : bool
        是否输出详细信息
    
    Returns:
    --------
    a : ndarray
        解向量 (复数形式)
    """
    
    # 构建正规方程 S @ a = b
    S = V.conj().T @ V  # V^H @ V
    b = V.conj().T @ L_rho  # V^H @ L_rho
    
    # 计算并输出矩阵的秩信息
    rank_S = np.linalg.matrix_rank(S)
    augmented = np.hstack((S, b.reshape(-1, 1)))
    rank_augmented = np.linalg.matrix_rank(augmented)
    
    if verbose:
        print(f"rank(S) = {rank_S}")
        print(f"rank([S | b]) = {rank_augmented}")
        
        # 输出条件数信息
        cond_S = np.linalg.cond(S)
        print(f"Condition number of S: {cond_S:.4e}")
        print("\n=== Iterative Refinement Process ===")
    
    # 检查系统的相容性
    if rank_S != rank_augmented:
        print("Warning: System is inconsistent (rank(S) ≠ rank([S|b]))")
        print("Using least squares solution...")
    
    # 步骤1: 获得初始解 (使用最小二乘法)
    try:
        a = lstsq(S, b, lapack_driver='gelsy')[0]
        if verbose:
            print("Initial solution computed using least squares")
    except Exception as e:
        print(f"Failed to compute initial solution: {e}")
        return np.zeros(S.shape[1], dtype=complex)
    
    # 初始残差
    initial_residual = np.linalg.norm(S @ a - b)
    if verbose:
        print(f"Initial residual norm: {initial_residual:.4e}")
    
    # 迭代精化过程
    for i in range(max_iter):
        # 步骤2: 计算当前残差 r = b - S @ a
        r = b - S @ a
        residual_norm = np.linalg.norm(r)
        
        if verbose:
            print(f"Iteration {i+1}: Residual norm = {residual_norm:.4e}")
        
        # 步骤3: 检查收敛
        if residual_norm < tolerance:
            if verbose:
                print(f"Converged after {i+1} iterations")
            break
        
        # 步骤4: 求解修正方程 S @ da = r
        try:
            da = lstsq(S, r, lapack_driver='gelsy')[0]
        except Exception as e:
            if verbose:
                print(f"Iteration {i+1} failed: {e}")
            break
        
        # 步骤5: 更新解 a = a + da
        a = a + da
        
        # 检查是否有改进，如果连续几次迭代没有改进则停止
        if i > 2 and residual_norm >= tolerance:
            prev_residuals = []
            for j in range(max(0, i-2), i+1):
                prev_residuals.append(residual_norm)
            
            # 如果最近3次迭代残差变化很小，可能已经达到数值极限
            if len(prev_residuals) >= 3:
                residual_changes = [abs(prev_residuals[k] - prev_residuals[k-1]) 
                                  for k in range(1, len(prev_residuals))]
                if all(change < tolerance * 10 for change in residual_changes):
                    if verbose:
                        print(f"Convergence stalled after {i+1} iterations")
                    break
    
    # 计算最终残差
    final_residual = np.linalg.norm(S @ a - b)
    
    if verbose:
        print(f"Final residual norm for S @ a - b: {final_residual:.4e}")
        
        # 输出改进情况
        improvement_factor = initial_residual / final_residual if final_residual > 0 else float('inf')
        print(f"Improvement factor: {improvement_factor:.2e}")
        
        # 检查解的性质
        max_imag = np.max(np.abs(np.imag(a)))
        if max_imag > 1e-12:
            print(f"Warning: Solution has significant imaginary parts: {max_imag:.4e}")
    
    # 输出必要的信息
    print(f"rank(S) = {rank_S}")
    print(f"rank([S | b]) = {rank_augmented}")
    print(f"Residual norm for S @ a - b: {final_residual:.4e}")
    
    return a

def iterative_refinement_solve_real_output(V, L_rho, max_iter=15, tolerance=1e-15, verbose=True):
    """
    迭代精化求解，返回实数解 (如果虚部可忽略)
    
    Parameters:
    -----------
    V : ndarray, shape (m, n)
        输入复数矩阵
    L_rho : ndarray, shape (m,)
        右端复数向量
    max_iter : int
        最大迭代次数
    tolerance : float
        收敛容差
    verbose : bool
        是否输出详细信息
    
    Returns:
    --------
    a : ndarray
        解向量 (如果虚部可忽略则返回实数，否则返回复数)
    """
    
    # 调用复数版本
    a_complex = iterative_refinement_solve(V, L_rho, max_iter, tolerance, verbose)
    
    # 检查是否可以转换为实数
    max_imag = np.max(np.abs(np.imag(a_complex)))
    
    if max_imag < 1e-12:  # 虚部可忽略
        if verbose:
            print("Converting to real solution (imaginary parts negligible)")
        return np.real(a_complex)
    else:
        if verbose:
            print(f"Keeping complex solution (max imaginary part: {max_imag:.4e})")
        return a_complex

# 简化版本 - 直接替换您原来的函数
def improved_iterative_solve(V, L_rho):
    """
    简化版迭代精化求解，直接替换您的原函数
    输出格式与您的需求完全一致
    """
    return iterative_refinement_solve_real_output(V, L_rho, max_iter=15, tolerance=1e-15, verbose=True)

# 使用示例
if __name__ == "__main__":
    # 示例用法：
    # 
    # # 加载您的数据
    # V = np.load(path+'V.npy').astype(np.complex128)
    # L_rho = np.load(path+'L_rho.npy').astype(np.complex128)
    # 
    # # 调用函数
    # solution = improved_iterative_solve(V, L_rho)
    # 
    # # 或者使用详细版本
    # solution = iterative_refinement_solve(V, L_rho, max_iter=20, tolerance=1e-16)
    
    pass